#!/usr/bin/env python
# -*- coding:UTF-8

"""
Created on 2017-09-26

@author: tangcheng
"""

import argparse
import logging
import logging.handlers
import os
import pwd
import re
import resource
import shutil
import subprocess
import sys
import time

g_version = "v1.0"
g_tag_line = """# ====== Add by cspgsetup ======"""


def run_cmd(cmd):
    logging.debug("Run: %s" % cmd)
    os.system(cmd)


def run_cmd_result(cmd):
    try:
        p = subprocess.Popen(
            cmd,
            shell=True,
            close_fds=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        out_msg = p.stdout.read()
        err_msg = p.stderr.read()
        err_code = p.wait()
        logging.debug("Run: %s, return %s" % (cmd, out_msg))
    except Exception as e:
        raise e

    return err_code, err_msg, out_msg


def modify_config_type1(config_file, modify_item_dict, deli_type=1, is_backup=True):
    """
    修改以等号或空格分隔的配置文件，默认为等号分隔的文件即（k = v）格式的配置文件，如postgresql.conf, /etc/sysctl.conf这些文件都是这种格式。
    :param config_file:
    :param modify_item_dict:
    :param deli_type: 1表示以等号分隔的配置文件，2表示以空格分隔的配置文件
    :param is_backup:
    :return:

    修改配置文件:
      1. 如果在文件中只有相应被注释掉的配置项，则新的配置项加在注释掉的配置项后面。
      2. 如果已存在配置项，则替换原有的配置项。
      3. 如果文件中不存在的配置项，则添加到文件尾
    例如modify_item_dict={'port':'5444'}，配置文件中只存在port的注释掉的配置：
      ...
      listen_addresses = '*'
      #port = 5432                            # (change requires restart)
      max_connections = 100                   # (change requires restart)
      ...

    执行后的结果是在此被注释掉的配置项后面加上新的配置项，结果变为：
      ...
      listen_addresses = '*'
      #port = 5432                            # (change requires restart)
      port = 5444
      max_connections = 100                   # (change requires restart)
      ...

    如果配置文件中存在port的注释掉的配置项和未被注释掉的相应的配置项：
      ...
      listen_addresses = '*'
      #port = 5432                            # (change requires restart)
      port = 5433
      max_connections = 100                   # (change requires restart)
      ...

    执行后的结果是在此被注释掉的配置项后面加上新的配置项，结果变为：
      ...
      listen_addresses = '*'
      #port = 5432                            # (change requires restart)
      port = 5444
      max_connections = 100                   # (change requires restart)
      ...

    :param config_file:
    :param modify_item_dict:
    :return:
    """

    fp = open(config_file, "r")
    ori_lines = fp.readlines()
    fp.close()

    # 下面的操作先找各个配置项的位置
    # item_line_num_dict1和item_line_num_dict2分别记录相应的配置项在文件中的行号。
    #   只是item_line_num_dict1字典中key是行号，而value是相应的配置项名称
    #   而item_line_num_dict2字典中key是配置项名称，而value是相应的行号
    item_line_num_dict1 = {}
    item_line_num_dict2 = {}

    # item_comment_line_num_dict1和item_comment_line_num_dict2分别记录配置文件中被注释掉的配置项在文件中的行号。
    item_comment_line_num_dict1 = {}
    item_comment_line_num_dict2 = {}

    i = 0
    for line in ori_lines:
        line = line.strip()
        if deli_type == 1:
            cells = line.split("=")
        else:
            cells = line.split()
        if len(cells) < 2:
            i += 1
            continue
        item_name = cells[0].strip()
        if item_name[0] == "#":
            if item_name[1:].strip() in modify_item_dict:
                item_comment_line_num_dict1[i] = item_name[1:]
                item_comment_line_num_dict2[item_name[1:]] = i
        if item_name in modify_item_dict:
            item_line_num_dict1[i] = item_name
            item_line_num_dict2[item_name] = i
        i += 1

    # 如果已存在相应的配置项，即使也存在注释掉的配置项，则就不能在已注释掉的配置项后再加上新配置项了，需要替换掉的已存在的配置项
    for item_name in item_comment_line_num_dict2:
        if item_name in item_line_num_dict2:
            i = item_comment_line_num_dict2[item_name]
            del item_comment_line_num_dict1[i]

    # 如果配置项在item_line_num_dict1中存在或在item_comment_line_num_dict1，则添加新配置项
    i = 0
    new_lines = []
    for line in ori_lines:
        line = line.strip()
        if i in item_line_num_dict1:
            if deli_type == 1:
                new_line = "%s = %s" % (
                    item_line_num_dict1[i],
                    modify_item_dict[item_line_num_dict1[i]],
                )
            else:
                new_line = "%s %s" % (
                    item_line_num_dict1[i],
                    modify_item_dict[item_line_num_dict1[i]],
                )
            new_lines.append(new_line)
        elif i in item_comment_line_num_dict1:
            # 如新行加到注释行的下一行处
            new_lines.append(line)
            item_name = item_comment_line_num_dict1[i]
            if deli_type == 1:
                new_line = "%s = %s" % (item_name, modify_item_dict[item_name])
            else:
                new_line = "%s %s" % (item_name, modify_item_dict[item_name])
            new_lines.append(new_line)
        else:
            new_lines.append(line)
        i += 1

    # 把配置文件中不存在的配置项，添加到文件尾（按item_name排序）
    items = sorted(modify_item_dict.keys())
    for item_name in items:
        if (
            item_name not in item_line_num_dict2
            and item_name not in item_comment_line_num_dict2
        ):
            if deli_type == 1:
                new_line = "%s = %s" % (item_name, modify_item_dict[item_name])
            else:
                new_line = "%s %s" % (item_name, modify_item_dict[item_name])
            new_lines.append(new_line)

    if is_backup:
        config_file_bak = "%s.%s" % (config_file, time.strftime("%Y%m%d%H%M%S"))
        shutil.copy(config_file, config_file_bak)

    fp = open(config_file, "w")
    content = "\n".join(new_lines)
    fp.write(content)
    fp.close()


def modify_config_type2(config_file, modify_item_dict, is_backup=True):
    """
    通过正则表达式匹配进行更新配置文件，如果匹配了，则替换，如果没有匹配上，会跳过，也不会在末尾添加。
    :param config_file:
    :param modify_item_dict: 这是一个字典，k为一个正则表达式，v为要替换的行
    :param is_backup:
    :return:
    """

    fp = open(config_file, "r")
    ori_lines = fp.readlines()
    fp.close()

    new_lines = []
    for line in ori_lines:
        line = line.strip()
        matched = False
        for k in modify_item_dict:
            if re.match(k, line):
                new_lines.append(modify_item_dict[k])
                matched = True
                break

        if not matched:
            new_lines.append(line)

    if is_backup:
        config_file_bak = "%s.%s" % (config_file, time.strftime("%Y%m%d%H%M%S"))
        shutil.copy(config_file, config_file_bak)

    fp = open(config_file, "w")
    content = "\n".join(new_lines)
    fp.write(content)
    fp.close()


def append_config_file(file_name, start_tag_line, append_contents):
    """
    这函数用于在指定的配置文件中添加内容，添加的内容的第一行中有start_tag_line，用于标识，我们添加的同容的开始
    :param file_name:
    :param start_tag_line: 这是一个标识字符串，用于从哪一行开始添加自定义的内容
    :param append_contents: 要添加的内容
    :return:
    """

    fp = open(file_name)
    content = fp.read()
    fp.close()
    lines = content.split("\n")
    flag = 0
    new_lines = []
    for line in lines:
        new_lines.append(line)
        if line == start_tag_line:
            flag = 1
            break

    # 第一次需要加入这个标志行，以后修改就只修改此标志行后面的内容
    if not flag:
        new_lines.append(start_tag_line)
    new_contents = "\n".join(new_lines)
    new_contents = "%s\n%s" % (new_contents, append_contents)

    fp = open(file_name, "w")
    fp.write(new_contents)
    fp.close()


def add_os_user(os_uid, os_user):
    logging.info("Begin add user and group ...")
    cmd = "groupadd -g %d %s" % (os_uid, os_user)
    run_cmd(cmd)
    cmd = "useradd -m -u %(uid)d -g %(uid)d -s /bin/bash %(user)s" % {
        "uid": os_uid,
        "user": os_user,
    }
    run_cmd(cmd)
    logging.info("Add user and group ok.")


def del_os_user(os_user):
    logging.info("Begin delete postgres user and group")
    cmd = "userdel -r %s" % os_user
    run_cmd(cmd)
    logging.info("Delete user(%s) and group." % os_user)


def get_memory_stat():
    mem = {}
    f = open("/proc/meminfo")
    lines = f.readlines()
    f.close()
    for line in lines:
        if len(line) < 2:
            continue
        name = line.split(":")[0]
        var = line.split(":")[1].split()[0]
        mem[name] = int(var) * 1024
    return mem


def init_log(level):
    """
    初使化日志，让日志即能屏幕输出，也可以输出到dycm.log中

    @param level: 日志级别，可以取的值为logging.INFO, logging.DEBUG等等值
    @return:
    """

    logger = logging.getLogger()

    formatter = logging.Formatter("%(asctime)s %(levelname)s %(message)s")

    stdout_handle = logging.StreamHandler()
    stdout_handle.setFormatter(formatter)
    logger.addHandler(stdout_handle)

    # log_file = os.path.realpath("/tmp/cspgsetup.log")
    # file_handle = logging.handlers.RotatingFileHandler(log_file, maxBytes=10 * 1024 * 1024, backupCount=5)
    # file_handle.setFormatter(formatter)
    # logging.addHandler(file_handle)
    logger.setLevel(level)


def sysctl_set():
    # 先设置共享内存的参数：

    mem = get_memory_stat()

    shmmax = mem["MemTotal"]
    shmall = shmmax / resource.getpagesize()
    ci_dict = {}

    ci_dict["kernel.shmall"] = shmall
    ci_dict["kernel.shmmax"] = shmmax

    # 信号量配置
    # SEMMSL(信号集的最大信号量数）：需要大于17，我们取250
    # SEMMNS（整个系统范围内的最大信息量数）：SEMMSL * SEMMNI = 650*250
    # SEMOPM（semop函数在一次调用中所能操作一个信号量集中最大的信号量数）: 与SEMMSL相同，为250
    # SEMMNI（信号量集的最大数目）：ceil(max_conntions + autovacuum_max_workers + 4)/16，如果对于10000个连接，而autovacuum_max_workers=3，所以SEMMNI最少要大于626，我们取650
    ci_dict["kernel.sem"] = "250 162500 250 650"

    # 虚拟内存配置
    ci_dict["vm.swappiness"] = 0
    ci_dict["vm.overcommit_memory"] = 2
    ci_dict["vm.overcommit_ratio"] = 50
    ci_dict["vm.dirty_background_ratio"] = 1
    ci_dict["vm.dirty_ratio"] = 2

    # kernel的其它配置参数
    ci_dict["kernel.sysrq"] = 1
    ci_dict["kernel.core_uses_pid"] = 1
    ci_dict["kernel.msgmnb"] = 65536
    ci_dict["kernel.msgmax"] = 65536
    ci_dict["kernel.msgmni"] = 2048

    ci_dict["net.ipv4.tcp_syncookies"] = 1
    ci_dict["net.ipv4.ip_forward"] = 0
    ci_dict["net.ipv4.conf.default.accept_source_route"] = 0
    ci_dict["net.ipv4.tcp_tw_recycle"] = 1
    ci_dict["net.ipv4.tcp_max_syn_backlog"] = 4096
    ci_dict["net.ipv4.conf.all.arp_filter"] = 1
    ci_dict["net.ipv4.ip_local_port_range"] = "1025 65535"
    ci_dict["net.core.netdev_max_backlog"] = 10000
    ci_dict["net.core.rmem_max"] = 2097152
    ci_dict["net.core.wmem_max"] = 2097152

    modify_config_type1("/etc/sysctl.conf", ci_dict)
    run_cmd("sysctl -p")


def ulimit_set():
    # 修改ulimit设置
    append_contents = (
        "* soft nofile 65536\n"
        "* hard nofile 65536\n"
        "* soft nproc 131072\n"
        "* hard nproc 131072"
    )
    append_config_file("/etc/security/limits.conf", g_tag_line, append_contents)

    # /etc/security/limits.d/20-nproc.conf
    nproc_conf = "/etc/security/limits.d/20-nproc.conf"
    if not os.path.exists(nproc_conf):
        # CentOS6.X没有20-nproc.conf文件，而是90-nproc.conf文件：
        nproc_conf = "/etc/security/limits.d/90-nproc.conf"
        if not os.path.exists(nproc_conf):
            logging.info(
                "Cant not find nproc file 20-nproc.conf or 90-nproc.conf in /etc/security/limits.d !!!"
            )
            return
    ci_dict = {"^\*\s+soft\s+nproc\s+\d+$": "*          soft    nproc     131072"}
    modify_config_type2(nproc_conf, ci_dict)


def sshd_config_set():
    # 修改sshd_config中的配置
    # ﻿/etc/ssh/sshd_config
    ci_dict = {"UseDNS": "no", "GSSAPIAuthentication": "no"}
    modify_config_type1("/etc/ssh/sshd_config", ci_dict, deli_type=2)


def init_os_env():
    # 修改sshd_config
    sshd_config_set()

    # 修改/etc/sysctl.conf
    sysctl_set()

    # ulimit配置
    ulimit_set()

    # 禁止SELinux
    ci_dict = {"﻿SELINUX": "disabled"}
    modify_config_type1("/etc/selinux/config", ci_dict)

    # 关闭防火墙
    run_cmd("systemctl disable firewalld")

    # io的默认调度策略改成deadline
    run_cmd('''grubby --update-kernel=ALL --args="elevator=deadline"''')

    # io的默认调度策略改成deadline
    run_cmd('''grubby --update-kernel=ALL --args="transparent_hugepage=never"''')


def init_options_parser():
    """
    :return:
    """

    usage = "%%prog %s [options]\n" % sys.argv[1]
    parser = argparse.ArgumentParser(usage=usage)
    parser.add_argument(
        "-l",
        "--loglevel",
        action="store",
        dest="loglevel",
        default="info",
        help="Specifies log level:  debug, info, warn, error, critical, default is info",
    )
    return parser


def pg_ver_to_int(pg_version):
    """
    :param pg_version:
    :return: 返回-1，表示是无效的version
    """

    cell = pg_version.split(".")
    if len(cell) == 2:
        return 1000000 * int(cell[0]) + 1000 * int(cell[1])
    elif len(cell) == 3:
        return 1000000 * int(cell[0]) + 1000 * int(cell[1]) + int(cell[2])
    else:
        return -1


def init_pg_db(pg_version):
    pg_install_dir = "/usr/pgsql-%s" % pg_version
    if not os.path.exists(pg_install_dir):
        sys.stderr.write(
            "PostgreSQL %s not install in %s" % (pg_version, pg_install_dir)
        )
        sys.exit(-1)

    # 先删除已存在的postgres用户
    del_os_user("postgres")

    # 建postgres用户
    add_os_user(701, "postgres")

    # 在.bash_profile中添加环境变量
    append_contents = (
        "export PGDATA=~/pgdata\n"
        "export PATH=%s/bin:$PATH\n"
        "export LD_LIBRARY_PATH=%s/lib:$LD_LIBRARY_PATH\n"
        "export PGHOST=/tmp\n"
        "export LANG=en_US.UTF-8" % (pg_install_dir, pg_install_dir)
    )
    append_config_file("/home/postgres/.bash_profile", g_tag_line, append_contents)

    # 建pgdata目录
    pgdata = "/home/postgres/pgdata"
    os.mkdir(pgdata, 0o700)
    upw = pwd.getpwnam("postgres")
    os.chown(pgdata, upw.pw_uid, upw.pw_gid)

    # 使用initdb初使化数据库
    run_cmd('''su - postgres -c "initdb"''')

    # 配置postgresql.conf文件
    ci_dict = {
        "listen_addresses": "'*'",
        "unix_socket_directories": "'/tmp'",
        "tcp_keepalives_idle": "5",
        "tcp_keepalives_interval": "5",
        "tcp_keepalives_count": "3",
        "maintenance_work_mem": "256MB",
        "bgwriter_delay": "50ms",
        "max_wal_senders": "10",
        "wal_keep_segments": "80",
        "hot_standby ": "on",
        "hot_standby_feedback": "on",
        "logging_collector": "on",
        "log_min_duration_statement": "10000",
        "track_io_timing": "on",
        "track_functions": "all",
        "autovacuum_vacuum_cost_delay": "1ms",
        "statement_timeout": "0",
        "lock_timeout": "60000",
    }

    if pg_ver_to_int(pg_version) < pg_ver_to_int("9.6"):
        ci_dict["wal_level"] = "hot_standby"
    else:
        ci_dict["wal_level"] = "replica"

    if pg_ver_to_int(pg_version) < pg_ver_to_int("9.5"):
        ci_dict["checkpoint_segments"] = "30"
    else:
        ci_dict["max_wal_size"] = "4GB"
        ci_dict["min_wal_size"] = "3GB"

    mem = get_memory_stat()
    # 总内存的1/4
    shared_mem = int(mem["MemTotal"]) / 1024 / 1024 / 4
    ci_dict["shared_buffers"] = "%dMB" % shared_mem
    if shared_mem < 1024:
        max_conn = 100
    else:
        max_conn = shared_mem / 8
    if max_conn > 5000:
        max_conn = 5000

    ci_dict["max_connections"] = "%d" % max_conn

    modify_config_type1("/home/postgres/pgdata/postgresql.conf", ci_dict)


def action_init_os_env():
    parser = init_options_parser()
    _options = parser.parse_args(sys.argv[1:])
    init_os_env()


def action_init_pg_db():
    parser = init_options_parser()
    parser.add_argument(
        "-y", "--yes", action="store_true", dest="yes", help="all select yes."
    )

    parser.add_argument(
        "-v",
        "--pgversion",
        action="store",
        dest="pgversion",
        default="9.6",
        help="Specifies the version of postgresql, example: 9.6、10、11、12， notice：/usr/pgsql-<version> must exists!",
    )

    options = parser.parse_args(sys.argv[1:])

    if not options.yes:
        is_continue = input(
            "Delete an existing database. All data will be lost. Do you want to continue?(y/n):\n"
        ).lower()
        if is_continue != "y":
            return
    # parser.print_help()
    init_pg_db(options.pgversion)


def main():
    init_log(logging.DEBUG)

    cmd_def_list = [
        [action_init_os_env, "init_os_env", "init os env."],
        [action_init_pg_db, "init_pg_db", "init a pg database."],
    ]

    prog_name = os.path.basename(sys.argv[0])
    usage = (
        "%s v%s Author: tangcheng\n"
        "usage: %s <command> [options]\n"
        "    command can be one of the following:" % (prog_name, g_version, prog_name)
    )

    for cmd_item in cmd_def_list:
        usage += "\n      %-16s: %s" % (cmd_item[1], cmd_item[2])

    cmd_name_list = []
    for cmd_item in cmd_def_list:
        cmd_name_list.append(cmd_item[1])

    if len(sys.argv) == 1:
        print(usage)
        sys.exit(0)

    for cmd_item in cmd_def_list:
        if sys.argv[1] == cmd_item[1]:
            print("=== %s V0.1 Author: osdba ===\n" % prog_name)
            cmd_item[0]()
            sys.exit(0)

    print(usage)
    sys.exit(1)


if __name__ == "__main__":
    main()
